import io
import os
from seal import *
from app.pkg.snark_verifier import SNARKVerifier
from app.pkg.utils import *

class FileEncryptor:
    def __init__(self, context, key_manager, oss_handler):
        self.context = context
        self.key_manager = key_manager
        self.oss_handler = oss_handler
        self._current_crs = None

    def encrypt_file(self, file_obj, filename, circuit):
        try:
            file_base, _ = os.path.splitext(filename)
            oss_dir = f"{file_base}/"
            secret_key, public_key = self.key_manager.generate_keys(self.context)
            secret_key_bytes = self.key_manager.serialize_key_to_bytes(secret_key)
            public_key_bytes = self.key_manager.serialize_key_to_bytes(public_key)
            # secret_path = self.key_manager.save_key(secret_key, 'private_key.key')
            # public_path = self.key_manager.save_key(public_key, 'public_key.key')
            self.key_manager.upload_private_key(secret_key_bytes, f'{oss_dir}keys/private_key.key')
            self.oss_handler.upload_file(public_key_bytes, f'{oss_dir}keys/public_key.key')

            file_bytes = file_obj.read()
            header, content = file_bytes[:512], file_bytes[512:]
            content_ints = list(content)

            encoder = BatchEncoder(self.context)
            encryptor = Encryptor(self.context, public_key)
            verifier = SNARKVerifier(self.context)

            if self._current_crs is None:
                self._current_crs, vrs = verifier.setup(circuit)
            crs, vrs = self._current_crs, verifier.vrs

            # 上传 crs
            crs_json = json.dumps({
                'h': verifier.h.coeffs.tolist(),
                'qrp': {
                    'V': [v.coeffs.tolist() for v in crs[1].V],
                    'W': [w.coeffs.tolist() for w in crs[1].W],
                    'Y': [y.coeffs.tolist() for y in crs[1].Y],
                    't': crs[1].t.coeffs.tolist()
                },
                'crs_data': crs[2]
            }, cls=NumpyEncoder)
            self.oss_handler.upload_file(io.BytesIO(crs_json.encode('utf-8')), oss_dir + 'evidence.crs')

            # 上传 vrs
            vrs_json = json.dumps(vrs[2], cls=NumpyEncoder)
            self.oss_handler.upload_file(io.BytesIO(vrs_json.encode('utf-8')), oss_dir + 'evidence.vrs')

            # 分块加密
            chunk_size = 8192
            chunk_info = []

            for i in range(0, len(content_ints), chunk_size):
                chunk = content_ints[i:i + chunk_size]
                original_chunk_length = len(chunk)

                if len(chunk) < chunk_size:
                    chunk += [0] * (chunk_size - len(chunk))

                chunk_np = np.array(chunk, dtype=np.int64).flatten()
                plaintext = encoder.encode(chunk_np)
                ciphertext = encryptor.encrypt(plaintext)
                compressed = verifier.compress_ciphertext(ciphertext)

                # mem_file = io.BytesIO()
                # compressed.save(mem_file)
                # mem_file.seek(0)

                mem_bytes = self.key_manager.serialize_key_to_bytes(compressed)

                chunk_name = f"chunk_{i}"
                oss_path = oss_dir + 'chunks/' + chunk_name
                self.oss_handler.upload_file(mem_bytes, oss_path)

                chunk_info.append({
                    "file": 'chunks/' + chunk_name,
                    "length": original_chunk_length
                })

            # 上传 chunk_info.json
            chunk_info_json = json.dumps(chunk_info)
            self.oss_handler.upload_file(io.BytesIO(chunk_info_json.encode('utf-8')), oss_dir + 'chunk_info.json')

            # 上传 header
            self.oss_handler.upload_file(io.BytesIO(header), oss_dir + 'header.bin')

        except Exception as e:
            print(f"[加密失败] Encryption error: {e}")
            raise